import gc
import math
import shutil
import time

import numpy as np
import cv2 as cv
import os
import matplotlib.pyplot as plt
from insects_detection.project_utils import project_tools
from math import pi


class batch_manager():
    def __init__(self, batch_target, batch_size):
        self.batch_size = batch_size
        self.batch_target = batch_target
        self.batch_start = 0
        self.batch_end = 0
        self.batch_num = math.floor(len(self.batch_target) / self.batch_size)

    def next_batch(self):
        if self.batch_num > len(self.batch_target):
            return self.batch_target
        else:
            self.batch_end = self.batch_start + self.batch_size
            batch_collection = self.batch_target[self.batch_start:self.batch_end]
            self.batch_start = self.batch_end
            return batch_collection

    def has_next_batch(self):
        if self.batch_end > len(self.batch_target):
            return False
        else:
            return True


class pic_tools():
    @staticmethod
    def rename_imgs(gallery_path, green_house_tag, img_type):
        if project_tools.system_platform == 'linux':
            filepath=[]
            for root, dirs, files in os.walk(gallery_path):
                for name in files:
                    filepath.append(os.path.join(root, name))
            filepath.sort()
            root = gallery_path+os.sep
            for each_path, index in zip(filepath,range(len(filepath))):
                row = int((index) / 8) + 1
                column = (math.floor(index / 2)) % 4 + 1
                side_mediate = (index) % 4 + 1
                side = "A" if side_mediate % 2 != 0 else "B"
                filename=each_path.rsplit(os.sep,1)[1]
                new_filename = green_house_tag + '-' + str(row) + "-" + str(column) + "-" + side
                print(filename + '===>' + new_filename)
                os.renames(each_path, os.path.join(root, new_filename + img_type))
        else:
            for root, dirs, files in os.walk(gallery_path):
                for name, index in zip(files, range(len(files))):
                    row = int((index) / 8) + 1
                    column = (math.floor(index / 2)) % 4 + 1
                    side_mediate = (index) % 4 + 1
                    side = "A" if side_mediate % 2 != 0 else "B"
                    # print(row,column,side)
                    filepath = (os.path.join(root, name))
                    filename = name.split(".")[0]
                    new_filename = green_house_tag + '-' + str(row) + "-" + str(column) + "-" + side
                    print(name + '===>' + new_filename)
                    os.renames(filepath, os.path.join(root, new_filename + img_type))

    @staticmethod
    def read_imgs(img_url):
        # 使用imdecode函数进行读取
        img = cv.imdecode(np.fromfile(img_url, dtype=np.uint8), -1)
        return img

    @staticmethod
    def read_imgs_from_dir(dir):
        filepath = []
        for root, dirs, files in os.walk(dir):
            for name in files:
                filepath.append(os.path.join(root, name))
        return filepath

    # 将处理后的图片进行存储
    @staticmethod
    def write_single_imgs(store_path, processed_img,
                          original_img_name=time.strftime("%Y%m%d%H%M%S", time.localtime()) + '.jpg'):
        # img_name=str(processed_img).rsplit(os.sep,1)[1]
        # cv.imwrite(store_path+os.sep+original_img_name, processed_img)
        cv.imencode('.jpg', processed_img)[1].tofile(store_path + os.sep + original_img_name)
        print(original_img_name + '===>' + store_path + os.sep + original_img_name + ' 写入成功!')
        return original_img_name, store_path + os.sep + original_img_name

    @staticmethod
    def write_all_processed_imgs(from_dir, to_dir, async_run=False, celery_object=None, enable_batch=False,batch_size=12):
        origin_imgs_path = pic_tools.read_imgs_from_dir(from_dir)
        sub_folder_path = project_tools.prepare_img_output_dir(from_dir, to_dir)
        if os.listdir(sub_folder_path):
            shutil.rmtree(sub_folder_path)
            os.mkdir(sub_folder_path)
            print('已覆盖 '+sub_folder_path)
        # deprecated:enable_batch——>分块处理图像（已弃用）
        if enable_batch:
            bm = batch_manager(origin_imgs_path, batch_size)
            while bm.has_next_batch():
                img_batch = bm.next_batch()
                for img in img_batch:
                    _, filename = os.path.split(img)
                    ori_img = pic_tools.read_imgs(img)
                    processed_img = pic_tools.cutout_imgs(ori_img)
                    pic_tools.write_single_imgs(sub_folder_path, processed_img, filename)
        # 使用celery task进行异步处理
        elif async_run:
            total = len(origin_imgs_path)
            message = "preparing to process original images,please wait..."
            while celery_object == None:
                celery_object.update_state(state='INITIATING',
                                           meta={'current': 0, 'total': total,
                                                 'status': message})
                time.sleep(0.5)
            finished = 0
            for ori_img_path in origin_imgs_path:
                _, filename = os.path.split(ori_img_path)
                ori_img = pic_tools.read_imgs(ori_img_path)
                processed_img = pic_tools.cutout_imgs(ori_img)
                pic_tools.write_single_imgs(sub_folder_path, processed_img, filename)
                finished += 1
                message = "Now processing " + filename + "~"
                celery_object.update_state(state='IMAGE_PREPROCESS_PROGRESS',
                                           meta={'current': finished, 'total': total,
                                                 'status': message})
            celery_object.update_state(state='IMAGE_PREPROCESS_COMPLETE',
                                       meta={'current': total, 'total': total,
                                             'status': "Succeed in writing processed images"})
        # 同步阻塞处理
        else:
            for ori_img_path in origin_imgs_path:
                _, filename = os.path.split(ori_img_path)
                ori_img = pic_tools.read_imgs(ori_img_path)
                processed_img = pic_tools.cutout_imgs(ori_img)
                pic_tools.write_single_imgs(sub_folder_path, processed_img, filename)
        return sub_folder_path

    @staticmethod
    def cutout_imgs(img):
        # 截取出粘虫板亮黄色部分
        lower = np.array([24, 43, 110], dtype="uint8")
        upper = np.array([124, 255, 255], dtype="uint8")
        hsv = cv.cvtColor(img, cv.COLOR_BGR2HSV)
        mask = cv.inRange(hsv, lower, upper)
        kernel = cv.getStructuringElement(cv.MORPH_RECT,(15,15))
        mask = cv.dilate(mask, kernel)
        image = cv.bitwise_and(img, img, mask=mask)
        gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)  # 转换为灰度图
        ret, dst = cv.threshold(gray, 90, 255, cv.THRESH_BINARY)  # 输入灰度图，输出二值图
        contours, hierarchy = cv.findContours(dst, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)  # 轮廓检测
        height, width = img.shape[:2]
        index = 0
        max = 0
        for c in range(len(contours)):
            x, y, w, h = cv.boundingRect(contours[c])
            if h > height or w > width:
                continue
            area = cv.contourArea(contours[c])
            if area > max:
                max = area
                index = c
        x, y, w, h = cv.boundingRect(contours[index])
        cropped = image[y:y + h - 50, x:x + w - 50]
        dst = cv.resize(cropped, (1500, 1200), interpolation=cv.INTER_CUBIC)
        return dst

    @staticmethod
    def prewitt_operator(img):
        grayImage = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
        # print(type(grayImage[0,0]))
        kernelx = np.array([[1, 1, 1], [0, 0, 0], [-1, -1, -1]], dtype=int)  # prewitt 1算子
        # print(type(kernelx[0,0]))
        kernely = np.array([[-1, 0, 1], [-1, 0, 1], [-1, 0, 1]], dtype=int)  # prewitt 2算子
        x = cv.filter2D(grayImage, cv.CV_16S, kernelx)
        y = cv.filter2D(grayImage, cv.CV_16S, kernely)
        # print(type(x[0,0]))
        absx = cv.convertScaleAbs(x)
        absy = cv.convertScaleAbs(y)
        # print(type(absx[0,0]))
        Prewitt = cv.addWeighted(absx, 0.5, absy, 0.5, 0)
        return Prewitt

    @staticmethod
    def getContours(Prewitt):
        ret, Prewitt = cv.threshold(Prewitt, 35, 255, cv.THRESH_BINARY)  # 输入灰度图，输出二值图
        element = cv.getStructuringElement(cv.MORPH_CROSS, (1, 1))  # 形态学去噪
        Prewitt = cv.morphologyEx(Prewitt, cv.MORPH_OPEN, element)  # 开运算去噪
        contours, hierarchy = cv.findContours(Prewitt, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
        cv.drawContours(Prewitt, contours, -1, (120, 0, 0), 2)  # 绘制轮廓
        return contours

    # 粘虫板虫害特征提取，包括9个颜色特征和4个形态学特征
    @staticmethod
    def extract_insect_characteristics(img):
        Prewitt = pic_tools.prewitt_operator(img)
        contours = pic_tools.getContours(Prewitt)
        count = 0
        margin = 0
        X_train = np.empty(shape=[0, 13])
        # 遍历找到的所有昆虫
        for cont in contours:
            ares = cv.contourArea(cont)  # 计算包围性状的面积
            if ares < 5 or ares > 120:  # 过滤面积小于10的形状
                continue
            rect = cv.minAreaRect(cont)
            box = np.int0(cv.boxPoints(rect))
            rect_w, rect_h = int(rect[1][0]) + 1, int(rect[1][1]) + 1
            h, w = img.shape[:2]
            if rect_w <= rect_h:
                x, y = box[1][0], box[1][1]
                M2 = cv.getRotationMatrix2D((x, y), rect[2], 1)
                rotated_image = cv.warpAffine(img, M2, (w, h))
                rotated_canvas = rotated_image[y:y + rect_h + margin + 1, x:x + rect_w + margin + 1]
            else:
                x, y = box[2][0], box[2][1]  # 旋转中心
                M2 = cv.getRotationMatrix2D((x, y), rect[2] + 90, 1)
                rotated_image = cv.warpAffine(img, M2, (w, h))
                rotated_canvas = rotated_image[y:y + rect_w + margin + 1, x:x + rect_h + margin + 1]
            if rotated_canvas.shape[0] == 0 or rotated_canvas.shape[1] == 0:
                continue
            Perimeter = cv.arcLength(cont, True)  # 计算目标啊的周长
            Complexity = (Perimeter ** 2 / (4 * pi * ares))  # 计算目标的复杂度
            dutycycle = ares / (rect[1][0] * rect[1][1])  # 计算害虫的占空比
            count += 1  # 总体计数加1
            b, g, r = cv.split(rotated_canvas)
            b_mean = np.mean(b)
            g_mean = np.mean(g)
            r_mean = np.mean(r)
            rotated_canvas_HSV = cv.cvtColor(rotated_canvas, cv.COLOR_BGR2HSV)
            h, s, v = cv.split(rotated_canvas_HSV)
            h_mean = np.mean(h)
            s_mean = np.mean(s)
            v_mean = np.mean(v)
            rotated_canvas_LAB = cv.cvtColor(rotated_canvas, cv.COLOR_BGR2LAB)
            L, A, B = cv.split(rotated_canvas_LAB)
            L_mean = np.mean(L)
            A_mean = np.mean(A)
            B_mean = np.mean(B)
            X_train = np.append(X_train, [
                [b_mean, g_mean, r_mean, h_mean, s_mean, v_mean, L_mean, A_mean, B_mean, ares, Perimeter, Complexity,
                 dutycycle]], axis=0)
        return X_train

    # 逐个抠取粘虫板上的昆虫图片并展示，人工在控制台输入昆虫标记
    @staticmethod
    def extract_single_insect_img(img, display_size=50):
        insects_tags = np.array([])
        Prewitt = pic_tools.prewitt_operator(img)
        contours = pic_tools.getContours(Prewitt)
        insects_contours = []  # 列表里存放满足面积要求的轮廓
        count = -1
        margin = 0
        insects_chas = np.empty(shape=[0, 13])
        is_labeled_right = True
        previous_insect_index = 0
        already_labbeled = 0
        # 遍历找到的所有昆虫
        for const in contours:
            ares = cv.contourArea(const)
            if ares >= 5 and ares <= 120:  # 过滤一些面积过小或者过大的区域
                insects_contours.append(const)
        # 获取每个昆虫的13项特征数据
        while count < len(insects_contours) - 1:
            count += 1
            if is_labeled_right:
                insect_index = np.random.randint(0, len(insects_contours))
            else:
                insect_index = previous_insect_index
                is_labeled_right = True
            # ares = cv.contourArea(insects_contours[insect_index])  # 计算包围性状的面积
            rect = cv.minAreaRect(insects_contours[insect_index])
            box = np.int0(cv.boxPoints(rect))
            rect_w, rect_h = int(rect[1][0]) + 1, int(rect[1][1]) + 1
            h, w = img.shape[:2]
            if rect_w <= rect_h:
                x, y = box[1][0], box[1][1]
                M2 = cv.getRotationMatrix2D((x, y), rect[2], 1)
                rotated_image = cv.warpAffine(img, M2, (w, h))
                rotated_canvas = rotated_image[y:y + rect_h + margin + 1, x:x + rect_w + margin + 1]
            else:
                x, y = box[2][0], box[2][1]  # 旋转中心
                M2 = cv.getRotationMatrix2D((x, y), rect[2] + 90, 1)
                rotated_image = cv.warpAffine(img, M2, (w, h))
                rotated_canvas = rotated_image[y:y + rect_w + margin + 1, x:x + rect_h + margin + 1]
            if rotated_canvas.shape[0] == 0 or rotated_canvas.shape[1] == 0:
                continue
            ares = cv.contourArea(insects_contours[insect_index])  # 计算选中符合要求的轮廓的面积
            Perimeter = cv.arcLength(insects_contours[insect_index], True)  # 计算目标啊的周长
            Complexity = (Perimeter ** 2 / (4 * pi * ares))  # 计算目标的复杂度
            dutycycle = ares / (rect[1][0] * rect[1][1])  # 计算害虫的占空比
            b, g, r = cv.split(rotated_canvas)
            b_mean = np.mean(b)
            g_mean = np.mean(g)
            r_mean = np.mean(r)
            rotated_canvas_HSV = cv.cvtColor(rotated_canvas, cv.COLOR_BGR2HSV)
            h, s, v = cv.split(rotated_canvas_HSV)
            h_mean = np.mean(h)
            s_mean = np.mean(s)
            v_mean = np.mean(v)
            rotated_canvas_LAB = cv.cvtColor(rotated_canvas, cv.COLOR_BGR2LAB)
            L, A, B = cv.split(rotated_canvas_LAB)
            L_mean = np.mean(L)
            A_mean = np.mean(A)
            B_mean = np.mean(B)
            rotated_canvas = cv.resize(rotated_canvas, (display_size, display_size), interpolation=cv.INTER_CUBIC)
            cv.destroyAllWindows()
            # cv.namedWindow('第'+str(insect_index+1)+'号昆虫', 2)
            cv.imshow('第' + str(insect_index + 1) + '号昆虫', rotated_canvas)
            cv.waitKey(1)
            tag = input('您当前已标记' + str(already_labbeled) + '只昆虫，请标记' + str(insect_index + 1) + '号昆虫:')
            if tag not in ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'q']:
                print('标记错误（仅支持0-9）,请重新标记！')
                is_labeled_right = False
                previous_insect_index = insect_index
                continue
            already_labbeled += 1
            if tag == 'q':
                return insects_chas, insects_tags
            else:
                insects_tags = np.append(insects_tags, int(tag))
            insects_chas = np.append(insects_chas, [
                [b_mean, g_mean, r_mean, h_mean, s_mean, v_mean, L_mean, A_mean, B_mean, ares, Perimeter, Complexity,
                 dutycycle]], axis=0)
        return insects_chas, insects_tags

    # 返回输入的昆虫标记以及该昆虫标记对应的特征数组
    @staticmethod
    def map_tags_to_insects_characteristics(insects_tags, img_dir):
        tag_num = len(insects_tags)
        filename = os.path.splitext(img_dir)[0].rsplit(os.sep, 1)[1]
        chas_data_dir = project_tools.project_root_path() + os.sep + project_tools.characteristic_dirname + os.sep + filename + '.txt'
        insects_chas = np.loadtxt(chas_data_dir)[0:tag_num]
        return insects_chas

    @staticmethod
    def save_random_extracted_insects_chas(insects_chas, filename):
        new_file_path = project_tools.project_root_path() + os.sep + project_tools.trainning_source_dirname + os.sep + filename + '.txt'
        with open(new_file_path, 'w'):
            pass
        try:
            np.savetxt(new_file_path, insects_chas, fmt=["%.8f"] * insects_chas.shape[1],
                       delimiter=" ", newline="\n")
            print('随机抽取昆虫特征保存成功==>' + filename + '.txt')
        except:
            print('随机抽取昆虫特征保存失败==>' + filename + '.txt')

    @staticmethod
    def save_insects_tags(insects_tags, filename):
        new_file_path = project_tools.project_root_path() + os.sep + project_tools.insects_tags_dirname + os.sep + filename + '.txt'
        with open(new_file_path, 'w'):
            pass
        try:
            np.savetxt(new_file_path, insects_tags, fmt=['%d'], newline="\n")
            print('昆虫标签保存成功==>' + filename + '.txt')
        except:
            print('昆虫标签保存失败==>' + filename + '.txt')

    @staticmethod
    def extract_all_imgs_characteristcs(img_dir, storage_dir, async_run=False, celery_object=None):
        imgs_path = pic_tools.read_imgs_from_dir(img_dir)
        sub_folder_path = project_tools.prepare_img_output_dir(img_dir, storage_dir)
        # celery异步非阻塞方式调用
        if async_run:
            total = len(imgs_path)
            finished = 0
            message = "Preparing to extract insects' chas,please wait..."
            while celery_object == None:
                celery_object.update_state(state='INITIATING',
                                           meta={'current': 0, 'total': total,
                                                 'status': message})
                time.sleep(0.5)
            for img_path in imgs_path:
                _, filename = os.path.split(img_path)
                fn_no_appendix = filename.split('.')[0]
                # with open(storage_dir + os.sep + fn_no_appendix + '.txt', 'w'):
                #     pass
                img = pic_tools.read_imgs(img_path)
                img_chas = pic_tools.extract_insect_characteristics(img)
                np.savetxt(sub_folder_path + os.sep + fn_no_appendix + '.txt', img_chas,
                           fmt=["%.8f"] * img_chas.shape[1],
                           delimiter=" ", newline="\n")
                finished += 1
                message = "Now identifying insects in image--->" + filename + ", it might take some time~"
                celery_object.update_state(state='CHARACTER_EXTRACTION_PROGRESS',
                                           meta={'current': finished, 'total': total,
                                                 'status': message})
                print(fn_no_appendix + " 数据特征已写入==> " + sub_folder_path + os.sep + fn_no_appendix + '.txt')
            celery_object.update_state(state='CHARACTER_EXTRACTION_COMPLETE',
                                       meta={'current': total, 'total': total,
                                             'status': "Succeed in identifying all insects, now generating result sheets~"})
        # 同步阻塞方式调用
        else:
            for img_path in imgs_path:
                _, filename = os.path.split(img_path)
                fn_no_appendix = filename.split('.')[0]
                # with open(sub_folder_path + os.sep + fn_no_appendix+'.txt', 'w'):
                #     pass
                img = pic_tools.read_imgs(img_path)
                img_chas = pic_tools.extract_insect_characteristics(img)
                np.savetxt(sub_folder_path + os.sep + fn_no_appendix + '.txt', img_chas,
                           fmt=["%.8f"] * img_chas.shape[1], delimiter=" ", newline="\n")
                print(fn_no_appendix + " 数据特征已写入==> " + sub_folder_path + os.sep + fn_no_appendix + '.txt')


class txt_tools():

    @staticmethod
    def read_txt():
        pass

    @staticmethod
    def read_files_from_dir(dir):
        filepath = []
        for root, dirs, files in os.walk(dir):
            for name in files:
                filepath.append(os.path.join(root, name))
        if project_tools.system_platform == 'linux':
            filepath.sort()
        return filepath
